import os
from datetime import timedelta
from typing import Union, Dict, Any

from asyncer import asyncify
from chainlit.data.storage_clients.base import BaseStorageClient
from minio import Minio, S3Error

from app.config.logging import logger
from app.decorator.timeit import timeit

make_async = asyncify

MINIO_CONFIG = {
    "endpoint": os.getenv("MINIO_ENDPOINT", "http://minio:9000"),
    "access_key": os.getenv("MINIO_ACCESS_KEY", "minio"),
    "secret_key": os.getenv("MINIO_SECRET_KEY", "minio"),
    "bucket_name": os.getenv("MINIO_BUCKET_NAME", "rag-bucket"),
    "storage_expiry_time": os.getenv("MINIO_STORAGE_EXPIRY_TIME", 3600),
}


class MinioStorageClient(BaseStorageClient):
    """
    使用Minio存储聊天记录信息
    参考S3Bucket
    """

    @timeit
    def __init__(self, endpoint: str,
                 access_key: str,
                 secret_key: str,
                 bucket_name: str,
                 secure: bool = False):
        self.endpoint = endpoint
        self.bucket_name = bucket_name
        self.client = Minio(endpoint, access_key, secret_key, secure=secure)
        self._ensure_bucket_exists()

    def sync_upload_file(self, object_key: str, data: Union[bytes, str], mime: str = "application/octet-stream",
                         overwrite: bool = True) -> Dict[str, Any]:
        """
        同步上传文件
        """
        try:
            self.client.put_object(bucket_name=self.bucket_name, object_name=object_key, data=data, content_type=mime,
                                   length=len(data))
            url = f"http://{self.endpoint}/{self.bucket_name}/{object_key}"
            return {"object_key": object_key, "url": url}
        except Exception as e:
            logger.warning(f"MinioStorageClient, upload_file error: {e}")
            return {}

    async def upload_file(self, object_key: str, data: Union[bytes, str], mime: str = "application/octet-stream",
                          overwrite: bool = True) -> Dict[str, Any]:
        """
        上传文件，主要用来保存会话记录
        """
        return await make_async(self.sync_upload_file)(object_key, data, mime=mime, overwrite=overwrite)

    def sync_delete_file(self, object_key: str) -> bool:
        """
        删除文件
        """
        try:
            self.client.remove_object(bucket_name=self.bucket_name, object_name=object_key)
            return True
        except Exception as e:
            logger.warning(f"MinioStorageClient, delete_file error: {e}")
            return False

    async def delete_file(self, object_key: str) -> bool:
        return await make_async(self.sync_delete_file)(object_key)

    def sync_get_read_url(self, object_key: str) -> str:
        """
        获取文件URL
        """
        try:
            # 生成预签名的URL
            url = self.client.get_presigned_url(
                "get_object",
                bucket_name=self.bucket_name,
                object_name=object_key,
                expires=timedelta(seconds=MINIO_CONFIG.get("storage_expiry_time", 3600)),
            )
            return url
        except Exception as e:
            logger.warning(f"MinioStorageClientMinioStorageClient, get_read_url error: {e}")
            return object_key

    async def get_read_url(self, object_key: str) -> str:
        return await make_async(self.sync_get_read_url)(object_key)

    @timeit
    def _ensure_bucket_exists(self):
        """确保存储桶存在"""
        try:
            if not self.client.bucket_exists(self.bucket_name):
                self.client.make_bucket(self.bucket_name)
        except S3Error as exc:
            raise RuntimeError(f"Failed to ensure bucket exists: {exc}")


# 初始化
minio_storage_client = MinioStorageClient(endpoint=MINIO_CONFIG["endpoint"],
                                          access_key=MINIO_CONFIG["access_key"],
                                          secret_key=MINIO_CONFIG["secret_key"],
                                          bucket_name=MINIO_CONFIG["bucket_name"])
